package com.shoulder.ssh.utils;

import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.shoulder.ssh.properties.SSHProperties;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;

@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ShellUtils extends SSHUtils {

    public ChannelShell initShell(Session session) throws JSchException {
        ChannelShell channelShell = (ChannelShell) session.openChannel("shell");
        channelShell.setPty(true); // 根据需要设置PTY，或者更详细地设置PTY类型和大小
        channelShell.connect();
        return channelShell;
    }


    public ChannelShell initShell(String ip, int port, String username, String password) throws JSchException {
        Session session = super.init(ip, port, username, password);
        ChannelShell channelShell = (ChannelShell) session.openChannel("shell");
        channelShell.setPty(true); // 根据需要设置PTY，或者更详细地设置PTY类型和大小
        channelShell.connect();
        return channelShell;
    }

    public void out(ChannelShell channelShell, TextMessage message) throws Exception {
        OutputStream out = channelShell.getOutputStream();
        out.write(message.getPayload().getBytes());
        out.flush();
    }


    public void out(ChannelShell channelShell, String message) throws Exception {
        OutputStream out = channelShell.getOutputStream();
        out.write(message.getBytes());
        out.flush();
    }


    public void send(ChannelShell channelShell, WebSocketSession session) throws Exception {

        try (InputStream in = channelShell.getInputStream()) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = in.read(buffer)) != -1) {
                String output = new String(buffer, 0, bytesRead);
                session.sendMessage(new TextMessage(output));
            }
        }

    }

    public void startForwarding(ChannelShell channel, WebSocketSession session) {
        new Thread(() -> {
            try (InputStream in = channel.getInputStream()) {
                byte[] buffer = new byte[1024];
                int bytesRead;
                while ((bytesRead = in.read(buffer)) != -1) {
                    session.sendMessage(new TextMessage(new String(buffer, 0, bytesRead, StandardCharsets.UTF_8)));
                }
            } catch (Exception e) {
                log.error("Shell forwarding error", e);
            }
        }).start();
    }

}
